import requests
import sqlite3
import json
import pandas as pd


# 初始化数据库并创建表格
def initialize_db():
    connection = sqlite3.connect('market_data.db')  # 创建或连接数据库
    db_cursor = connection.cursor()
    # 创建存储股票数据的表格
    db_cursor.execute('''
        CREATE TABLE IF NOT EXISTS stock_info (
            symbol TEXT, 
            name TEXT, 
            latest_price REAL, 
            price_change_rate REAL, 
            price_change_amount REAL, 
            trade_volume INTEGER, 
            trade_value REAL, 
            amplitude REAL, 
            high_price REAL, 
            low_price REAL, 
            opening_price REAL, 
            closing_price REAL
        )
    ''')
    connection.commit()
    return connection, db_cursor


# 插入股票数据
def save_data(db_cursor, stock_info):
    db_cursor.execute('''
        INSERT INTO stock_info (symbol, name, latest_price, price_change_rate, price_change_amount, trade_volume, trade_value, amplitude, high_price, low_price, opening_price, closing_price)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    ''', stock_info)


# 获取股票市场信息
def fetch_market_data(api_url):
    try:
        result = requests.get(api_url)
        if result.status_code == 200:
            data_start = result.text.find('(') + 1
            data_end = result.text.rfind(')')
            json_data = result.text[data_start:data_end]
            parsed_data = json.loads(json_data)
            return parsed_data['data']['diff']
        else:
            print(f"获取数据失败，状态码: {result.status_code}")
            return []
    except Exception as error:
        print(f"请求出错: {error}")
        return []


# 显示带有列名的数据并导出到Excel
def show_stock_data():
    connection = sqlite3.connect('market_data.db')
    query = "SELECT * FROM stock_info"
    # 使用 pandas 读取数据库中的数据
    dataframe = pd.read_sql_query(query, connection)

    # 显示数据
    print(dataframe)

    # 将数据导出到 Excel
    excel_filename = 'stock_data.xlsx'
    dataframe.to_excel(excel_filename, index=False)
    print(f"数据已囚禁到 {excel_filename} 文件中。")

    connection.close()


# 主逻辑：抓取并存储股票数据
def main_program():
    # 初始化数据库
    connection, db_cursor = initialize_db()

    # 获取用户输入的页码范围
    start_pg = int(input("你想从哪页开始: "))
    end_pg = int(input("到哪页结束: "))

    # 循环抓取指定页码的数据
    for page_num in range(start_pg, end_pg + 1):
        print(f"快了！已经到第 {page_num} 页了！")

        # 动态生成URL
        api_url = f"https://1.push2.eastmoney.com/api/qt/clist/get?cb=jQuery112402758990905568719_1728977574693&pn={page_num}&pz=20&po=1&np=1&ut=bd1d9ddb04089700cf9c27f6f7426281&fltt=2&invt=2&dect=1&wbp2u=|0|0|0|web&fid=f3&fs=m:0+t:6,m:0+t:80,m:1+t:2,m:1+t:23,m:0+t:81+s:2048&fields=f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f12,f13,f14,f15,f16,f17,f18,f20,f21,f23,f24,f25,f22,f11,f62,f128,f136,f115,f152&_=1728977574694"

        # 获取股票数据
        stock_list = fetch_market_data(api_url)

        # 遍历并存储每只股票的数据
        for stock in stock_list:
            stock_entry = (
                stock.get('f12'),  # 股票代码
                stock.get('f14'),  # 股票名称
                stock.get('f2'),  # 最新报价
                stock.get('f3'),  # 涨跌幅
                stock.get('f4'),  # 涨跌额
                stock.get('f5'),  # 成交量
                stock.get('f6'),  # 成交额
                stock.get('f7'),  # 振幅
                stock.get('f15'),  # 最高价
                stock.get('f16'),  # 最低价
                stock.get('f17'),  # 今开价
                stock.get('f18'),  # 昨收价
            )
            # 保存到数据库
            save_data(db_cursor, stock_entry)

    # 提交事务并关闭数据库
    connection.commit()
    connection.close()
    print("数据已成功逮捕并关押在数据库！")

    # 显示存储的数据并导出到 Excel
    show_stock_data()


# 运行程序
if __name__ == "__main__":
    main_program()
